# -*- coding: utf-8 -*-
# @Function: 最终测试脚本
# @Description: 读取测试问题文件，进行RAG流程测试，输出测试结果到Excel
# @Usage: 直接运行此脚本进行完整的测试流程
# @Dependencies: Z_config.py, Z_db_manager.py, Z_tokenizer.py, docx, openpyxl, langchain

import os
import sys
import re
import json
import pandas as pd
from typing import List, Dict, Tuple
from pathlib import Path
from docx import Document
from docx.text.paragraph import Paragraph as DocxParagraph
from docx.table import Table as DocxTable
from docx.oxml.text.paragraph import CT_P
from docx.oxml.table import CT_Tbl
from dotenv import load_dotenv
from langchain_core.prompts import PromptTemplate
from langchain_openai import ChatOpenAI

# 导入项目模块（使用新的包结构）
from D_db.core.config import DB_CONFIG, RESERVOIR_MAPPING
from D_db.core.db_manager import DBManager
from D_db.core.tokenizer import Tokenizer
from D_db.tools.prompt_utils import build_llm_context_from_qa

# 加载环境变量
load_dotenv()

# 阿里云百炼配置
DASHSCOPE_CONFIG = {
    "api_key": os.getenv('DASHSCOPE_API_KEY'),
    "base_url": os.getenv('BASE_URL'),
    "model": "qwen-long"
}

class TestFinalProcessor:
    def __init__(self):
        """初始化测试处理器"""
        self.db_manager = DBManager()
        self.tokenizer = Tokenizer()
        
        # 初始化阿里云百炼模型
        if DASHSCOPE_CONFIG["api_key"]:
            self.llm = ChatOpenAI(
                api_key=DASHSCOPE_CONFIG["api_key"],
                base_url=DASHSCOPE_CONFIG["base_url"],
                model=DASHSCOPE_CONFIG["model"],
                temperature=0.7
            )
            
            # 创建回答生成的提示模板
            self.answer_template = PromptTemplate(
                input_variables=["question", "context"],
                template="""你是一个专业的水库调度规程问答助手。请根据以下参考信息，回答用户的问题。

用户问题：
{question}

参考信息：
{context}

要求：
1. 基于参考信息提供准确、完整的回答
2. 如果参考信息不足以回答问题，请明确指出
3. 回答要简洁明了，直接针对问题给出答案
4. 不要在回答中提及"根据参考信息"等字样
5. 使用专业、客观的语气

请提供你的回答：
"""
            )
            
            # 创建回答生成链 (与 llm_query_tester 保持一致)
            self.answer_chain = self.answer_template | self.llm
            self.llm_available = True
        else:
            print("警告: 未设置DASHSCOPE_API_KEY，将跳过LLM回答生成")
            self.llm_available = False
    
    def extract_questions_from_docx(self, file_path: str) -> List[Dict[str, str]]:
        """从Word文档中提取问题和答案
        
        Args:
            file_path (str): Word文档路径
            
        Returns:
            List[Dict[str, str]]: 问题和答案列表
        """
        try:
            doc = Document(file_path)
            qa_pairs = []
            current_question = ""
            current_answer = ""
            is_answer_section = False

            # 统一的问答起始正则（更健壮）：
            # 支持行首可选序号（如“98.”、“104、”），支持【】或[]包裹，支持“问题/提问/问”后跟可选数字与可选冒号
            question_start_re = re.compile(
                r'^(?:\d+[.、]?\s*)?(?:[【\[]?\s*(?:问题|提问|问)\d*\s*[】\]]?\s*[:：]?)'
            )
            # 支持【】或[]包裹，支持“答案/回答/解答/答”与可选冒号
            answer_start_re = re.compile(
                r'^(?:[【\[]?\s*(?:答案|回答|解答|答)\s*[】\]]?\s*[:：]?)'
            )

            # 遍历文档中的段落和表格，保持原始顺序
            def iter_block_items(document: Document):
                parent_elm = document.element.body
                for child in parent_elm.iterchildren():
                    if isinstance(child, CT_P):
                        yield DocxParagraph(child, document)
                    elif isinstance(child, CT_Tbl):
                        yield DocxTable(child, document)

            for block in iter_block_items(doc):
                # 处理段落
                if isinstance(block, DocxParagraph):
                    text = block.text.strip()
                    if not text:
                        continue

                    # 检查是否是问题行（支持多种格式）
                    # 兼容：行首序号、【】/[] 包裹、“问题/提问/问”+可选数字+可选冒号
                    if question_start_re.search(text):
                        # 保存上一个问答对
                        if current_question and current_answer:
                            qa_pairs.append({
                                'question': current_question.strip(),
                                'answer': current_answer.strip()
                            })

                        # 提取问题内容，去除序号、中括号内容等
                        # 去除行首序号
                        question_text = re.sub(r'^\d+[.、]?\s*', '', text)
                        # 去除“问题/提问/问”起始标识（含可选数字、括号、冒号）
                        question_text = re.sub(r'^(?:[【\[]?\s*(?:问题|提问|问)\d*\s*[】\]]?\s*[:：]?\s*)', '', question_text)
                        question_text = re.sub(r'\[.*?\]', '', question_text)  # 去除方括号内容
                        question_text = re.sub(r'【.*?】', '', question_text)  # 去除中文方括号内容
                        current_question = question_text.strip()
                        current_answer = ""
                        is_answer_section = False

                    # 检查是否是答案行
                    elif answer_start_re.search(text):
                        is_answer_section = True
                        # 去除“答案/回答/解答/答”起始标识（含括号、冒号）
                        answer_text = re.sub(r'^(?:[【\[]?\s*(?:答案|回答|解答|答)\s*[】\]]?\s*[:：]?\s*)', '', text)
                        # 清除答案中的【】内容
                        answer_text = re.sub(r'\[.*?\]', '', answer_text)
                        answer_text = re.sub(r'【.*?】', '', answer_text)
                        current_answer = answer_text.strip()

                    # 如果在答案部分，继续收集答案内容
                    elif is_answer_section and current_question:
                        # 清除答案内容中的【】
                        clean_text = re.sub(r'\[.*?\]', '', text)
                        clean_text = re.sub(r'【.*?】', '', clean_text)
                        current_answer += " " + clean_text

                    # 如果没有明确的问题/答案标识，尝试智能识别
                    elif current_question and not is_answer_section:
                        # 如果当前有问题但还没有答案，且这行不像是新问题，则作为答案
                        if not re.search(r'^\d+[.、]', text) and not question_start_re.search(text):
                            # 清除内容中的【】
                            clean_text = re.sub(r'\[.*?\]', '', text)
                            clean_text = re.sub(r'【.*?】', '', clean_text)
                            current_answer += " " + clean_text
                            is_answer_section = True

                # 处理表格：将表格内容并入答案
                elif isinstance(block, DocxTable):
                    table_lines = []
                    for row in block.rows:
                        cells = [re.sub(r'【.*?】|\[.*?\]', '', (cell.text or '').strip()) for cell in row.cells]
                        line = " | ".join(cells).strip()
                        if line:
                            table_lines.append(line)
                    table_text = "\n".join(table_lines)

                    if table_text:
                        if current_question:
                            if current_answer:
                                current_answer += "\n" + table_text
                            else:
                                current_answer = table_text
                            is_answer_section = True

            # 保存最后一个问答对
            if current_question and current_answer:
                qa_pairs.append({
                    'question': current_question.strip(),
                    'answer': current_answer.strip()
                })

            return qa_pairs

        except Exception as e:
            print(f"读取文档 {file_path} 失败: {e}")
            return []
    
    def get_category_id_from_filename(self, filename: str) -> str:
        """根据文件名获取category_id
        
        Args:
            filename (str): 文件名
            
        Returns:
            str: category_id
        """
        # 去除文件扩展名
        name_without_ext = os.path.splitext(filename)[0]
        
        # 在RESERVOIR_MAPPING中查找匹配的category_id
        for reservoir_name, category_id in RESERVOIR_MAPPING.items():
            if reservoir_name in name_without_ext or name_without_ext in reservoir_name:
                return category_id
        
        print(f"警告: 未找到文件 {filename} 对应的category_id")
        return None
    
    def query_database(self, question: str, category_id: str, limit: int = 20) -> Tuple[List[Tuple], List[Tuple]]:
        """查询数据库获取相关问答对和表格数据
        
        Args:
            question (str): 问题
            category_id (str): 类别ID
            limit (int): 返回结果数量限制
            
        Returns:
            Tuple[List[Tuple], List[Tuple]]: (问答对查询结果, 表格查询结果)
        """
        qa_results = []
        table_results = []
        
        try:
            # 连接数据库
            self.db_manager.connect()
            
            # 对问题进行分词
            token_data = self.tokenizer.tokenize_question(question)
            tokens = token_data.get('tokens', [])
            
            if not tokens:
                print(f"问题分词结果为空: {question}")
                return qa_results, table_results
            
            # 查询问答对数据
            qa_results = self.db_manager.query_qa(tokens, limit, category_id)
            
            # 当前DBManager未实现表格数据查询，返回空列表占位
            table_results = []
            
            return qa_results, table_results
            
        except Exception as e:
            print(f"数据库查询失败: {e}")
            return qa_results, table_results
        finally:
            self.db_manager.close()
    
    def generate_llm_answer(self, question: str, qa_results: List[Tuple], table_results: List[Tuple]) -> Tuple[str, str]:
        """使用LLM生成回答
        
        Args:
            question (str): 问题
            qa_results (List[Tuple]): 问答对查询结果
            table_results (List[Tuple]): 表格查询结果
            
        Returns:
            Tuple[str, str]: (LLM生成的回答, 传给大模型的上下文全文)
        """
        if not self.llm_available:
            # 即使LLM不可用，也构建上下文以写入Excel
            context_for_llm, _ = build_llm_context_from_qa(qa_results) if qa_results else ("", [])
            return "LLM服务不可用", context_for_llm
        
        if not qa_results and not table_results:
            return "没有找到相关信息", ""
        
        try:
            # 使用共享模块构建传给LLM的上下文（问答对 + 去重上下文，忽略 row_data）
            context_for_llm, _ = build_llm_context_from_qa(qa_results)

            # 使用与 llm_query_tester 相同的链调用方式
            response = self.answer_chain.invoke({"question": question, "context": context_for_llm})
            answer = response.content if hasattr(response, 'content') else str(response)
            
            # 确保answer是字符串类型
            answer = str(answer) if answer is not None else ""
            
            # 清除生成答案中的【】内容
            answer = re.sub(r'\[.*?\]', '', answer)
            answer = re.sub(r'【.*?】', '', answer)
            
            return answer.strip(), context_for_llm
            
        except Exception as e:
            print(f"LLM回答生成失败: {e}")
            return f"回答生成失败: {str(e)}", ""
    
    def format_db_results(self, qa_results: List[Tuple], table_results: List[Tuple]) -> str:
        """格式化数据库查询结果
        
        Args:
            qa_results (List[Tuple]): 问答对查询结果
            table_results (List[Tuple]): 表格查询结果
            
        Returns:
            str: 格式化后的结果字符串
        """
        formatted_results = []
        
        # 格式化问答对结果
        if qa_results:
            formatted_results.append("问答对匹配结果：")
            for i, result in enumerate(qa_results[:5], 1):
                question_text = result[1] if len(result) > 1 else ""
                answer_text = result[2] if len(result) > 2 else ""
                
                # 确保是字符串类型
                question_text = str(question_text) if question_text is not None else ""
                answer_text = str(answer_text) if answer_text is not None else ""
                
                # 清除【】内容
                question_text = re.sub(r'\[.*?\]', '', question_text)
                question_text = re.sub(r'【.*?】', '', question_text)
                answer_text = re.sub(r'\[.*?\]', '', answer_text)
                answer_text = re.sub(r'【.*?】', '', answer_text)
                
                formatted_results.append(f"{i}. 问题：{question_text}")
                formatted_results.append(f"   答案：{answer_text}")
        
        # 格式化表格结果
        if table_results:
            if qa_results:  # 如果已有问答对结果，添加分隔
                formatted_results.append("")
            formatted_results.append("表格数据匹配结果：")
            for i, result in enumerate(table_results[:3], 1):
                keywords = result[0] if len(result) > 0 else ""
                table_data = result[1] if len(result) > 1 else ""
                
                # 确保是字符串类型
                keywords = str(keywords) if keywords is not None else ""
                table_data = str(table_data) if table_data is not None else ""
                
                formatted_results.append(f"{i}. 关键词：{keywords}")
                formatted_results.append(f"   数据：{table_data}")
        
        if not formatted_results:
            return "未找到匹配的数据"
        
        return "\n".join(formatted_results)
    
    def process_test_file(self, file_path: str) -> List[Dict]:
        """处理单个测试文件
        
        Args:
            file_path (str): 测试文件路径
            
        Returns:
            List[Dict]: 测试结果列表
        """
        filename = os.path.basename(file_path)
        print(f"\n处理文件: {filename}")
        
        # 获取category_id
        category_id = self.get_category_id_from_filename(filename)
        if not category_id:
            print(f"跳过文件 {filename}：无法确定category_id")
            return []
        
        print(f"匹配的category_id: {category_id}")
        
        # 提取问答对
        qa_pairs = self.extract_questions_from_docx(file_path)
        if not qa_pairs:
            print(f"文件 {filename} 中未找到问答对")
            return []
        
        print(f"提取到 {len(qa_pairs)} 个问答对")
        
        results = []
        for i, qa_pair in enumerate(qa_pairs, 1):
            question = qa_pair['question']
            original_answer = qa_pair['answer']
            
            print(f"\n处理问题 {i}: {question[:50]}...")
            
            # 初始化结果字典，确保序号和问题内容始终被记录
            result = {
                '序号': i,
                '问题内容': question,
                '原始答案': original_answer,
                '大模型答案': '测试失败',
                '数据库匹配结果': '测试失败'
            }
            
            try:
                # 查询数据库（同时查询问答对和表格数据）
                qa_results, table_results = self.query_database(question, category_id, 20)
                
                # 生成LLM回答，并获取传给大模型的上下文
                llm_answer, llm_prompt_context = self.generate_llm_answer(question, qa_results, table_results)

                # 更新成功的结果：将传给大模型的内容写入“数据库匹配结果”
                result.update({
                    '大模型答案': llm_answer,
                    '数据库匹配结果': llm_prompt_context
                })
                
            except Exception as e:
                print(f"处理问题 {i} 时发生错误: {e}")
                # 保留默认的失败信息，确保问题和序号被记录
                result.update({
                    '大模型答案': f'测试失败: {str(e)}',
                    '数据库匹配结果': f'测试失败: {str(e)}'
                })
            
            results.append(result)
        
        return results
    
    def save_results_to_excel(self, results: List[Dict], output_file: str):
        """保存结果到Excel文件
        
        Args:
            results (List[Dict]): 测试结果
            output_file (str): 输出文件路径
        """
        try:
            # 确保输出目录存在
            os.makedirs(os.path.dirname(output_file), exist_ok=True)
            
            # 创建DataFrame
            df = pd.DataFrame(results)
            
            # 保存到Excel
            with pd.ExcelWriter(output_file, engine='openpyxl') as writer:
                df.to_excel(writer, index=False, sheet_name='测试结果')
                
                # 调整列宽
                worksheet = writer.sheets['测试结果']
                for column in worksheet.columns:
                    max_length = 0
                    column_letter = column[0].column_letter
                    for cell in column:
                        try:
                            if len(str(cell.value)) > max_length:
                                max_length = len(str(cell.value))
                        except:
                            pass
                    adjusted_width = min(max_length + 2, 50)  # 最大宽度50
                    worksheet.column_dimensions[column_letter].width = adjusted_width
            
            print(f"结果已保存到: {output_file}")
            
        except Exception as e:
            print(f"保存结果失败: {e}")
    
    def run_tests(self):
        """运行所有测试"""
        print("开始执行测试流程...")
        
        # 测试问题文件夹路径
        test_folder = "测试问题"
        results_folder = "测试结果"
        
        if not os.path.exists(test_folder):
            print(f"错误: 测试问题文件夹 {test_folder} 不存在")
            return
        
        # 创建结果文件夹
        os.makedirs(results_folder, exist_ok=True)
        
        # 遍历测试文件
        test_files = [f for f in os.listdir(test_folder) if f.endswith(('.docx', '.doc'))]
        
        if not test_files:
            print(f"在 {test_folder} 文件夹中未找到Word文档")
            return
        
        print(f"找到 {len(test_files)} 个测试文件")
        
        for test_file in test_files:
            file_path = os.path.join(test_folder, test_file)
            
            # 处理测试文件
            results = self.process_test_file(file_path)
            
            if results:
                # 生成输出文件名
                output_filename = os.path.splitext(test_file)[0] + ".xlsx"
                output_path = os.path.join(results_folder, output_filename)
                
                # 保存结果
                self.save_results_to_excel(results, output_path)
            else:
                print(f"文件 {test_file} 处理失败或无结果")
        
        print("\n所有测试完成！")

def main():
    """主函数"""
    print("=== 最终测试脚本 ===")
    print("功能: 读取测试问题，进行RAG流程测试，输出结果到Excel")
    
    # 检查环境变量
    if not DASHSCOPE_CONFIG["api_key"]:
        print("警告: 未设置DASHSCOPE_API_KEY环境变量，LLM功能将不可用")
        print("请在.env文件中设置DASHSCOPE_API_KEY和BASE_URL")
    
    # 创建处理器并运行测试
    processor = TestFinalProcessor()
    processor.run_tests()

if __name__ == "__main__":
    main()